import sys
import time
import shutil
import subprocess
import hashlib
import ctypes
import os
import pefile
import psutil
from pathlib import Path
from typing import Optional


# === 日志初始化 ===
def init_logger():
    log_dir = Path(sys.executable).parent / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / "updater.log"
    import logging
    logger = logging.getLogger("Updater")
    logger.setLevel(logging.INFO)
    if logger.handlers:
        return logger
    formatter = logging.Formatter('%(asctime)s - %(levelname)s - [Updater] %(message)s')
    file_handler = logging.FileHandler(log_file, encoding='utf-8')
    file_handler.setFormatter(formatter)
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    return logger


logger = init_logger()

"""
    更新逻辑：
    1. 获取新版本信息
    2. 验证新版本文件
    3. 创建更新目录
    4. 移动旧文件到更新目录
    5. 移动新版本文件到目标目录
    6. 删除旧文件
    7. 启动新版本
"""
# === 共享工具函数 ===
def verify_sha256(file_path: Path, expected_hash: str) -> bool:
    if not file_path.exists():
        logger.error(f"[校验] 文件不存在：{file_path}")
        return False
    sha256 = hashlib.sha256()
    try:
        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha256.update(chunk)
        actual = sha256.hexdigest().lower()
        expected = expected_hash.lower()
        return actual == expected
    except Exception as e:
        logger.error(f"[校验] 校验失败: {e}", exc_info=True)
        return False


def is_64bit_file(file_path: Path) -> Optional[bool]:
    if not file_path.exists() or not file_path.is_file():
        return None
    try:
        pe = pefile.PE(str(file_path))
        if pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_AMD64"]:
            return True
        elif pe.FILE_HEADER.Machine == pefile.MACHINE_TYPE["IMAGE_FILE_MACHINE_I386"]:
            return False
        else:
            return None
    except:
        return None


def is_admin() -> bool:
    try:
        return ctypes.windll.shell32.IsUserAnAdmin()
    except:
        return False


def kill_process_by_name(process_name: str) -> bool:
    try:
        killed = False
        current_pid = os.getpid()
        for proc in psutil.process_iter(['name', 'pid']):
            if proc.info['name'] == process_name and proc.info['pid'] != current_pid:
                proc.kill()
                killed = True
        time.sleep(0.5)
        return killed
    except Exception as e:
        logger.error(f"[进程清理] 关闭进程失败: {e}", exc_info=True)
        return False


def force_delete(path: Path) -> bool:
    if not path.exists():
        return True
    try:
        if path.is_file():
            path.unlink(missing_ok=True)
            if path.exists():
                delete_path = str(path) + ".delete_on_reboot"
                os.rename(path, delete_path)
                return False
            return True
        else:
            shutil.rmtree(path, ignore_errors=True)
            return not path.exists()
    except Exception as e:
        logger.error(f"[删除] {path} 删除失败: {e}", exc_info=True)
        return False


def copy_with_retry(src: Path, dst: Path, max_retries: int = 3) -> bool:
    for i in range(max_retries):
        try:
            dst.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src, dst)
            if dst.exists() and dst.stat().st_size == src.stat().st_size:
                return True
        except Exception as e:
            logger.warning(f"[复制] 第 {i+1} 次失败: {e}")
            if i < max_retries - 1:
                time.sleep(1)
    return False


def handle_pending_replace(src: Path, dst: Path, description: str) -> bool:
    try:
        temp_dir = Path(sys.executable).parent / "update_pending"
        temp_dir.mkdir(exist_ok=True)
        temp_file = temp_dir / f"temp_{dst.name}"
        if not copy_with_retry(src, temp_file):
            return False
        batch_script = temp_dir / f"replace_{dst.stem}.bat"
        with open(batch_script, "w", encoding="utf-8") as f:
            f.write("@echo off\n")
            f.write(f"echo 正在替换 {description}...\n")
            f.write("timeout /t 3 /nobreak >nul\n")
            f.write(f'del /f /q "{dst}" 2>nul\n')
            f.write(f'copy /y "{temp_file}" "{dst}" 2>nul\n')
            f.write(f'del /f /q "{batch_script}"\n')
            f.write("exit\n")
        subprocess.Popen(str(batch_script), shell=True, close_fds=True, creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
        return True
    except Exception as e:
        logger.error(f"[延迟替换] {description} 失败: {e}", exc_info=True)
        return False


def restart_application(target_exe: Path):
    try:
        subprocess.Popen(f'cmd /c "start "" "{target_exe}" "', shell=True, close_fds=True, creationflags=subprocess.CREATE_NEW_PROCESS_GROUP)
    except Exception as e:
        logger.error(f"[重启] 失败: {e}", exc_info=True)
        ctypes.windll.user32.MessageBoxW(0, f"更新完成，请手动启动：{target_exe}", "更新完成", 0x40)


def main():
    logger.info("===== 更新器启动 =====")
    if len(sys.argv) != 5:
        error_msg = "参数错误！格式: updater.exe <DataCollector.exe> <updater.exe> <version.txt> <SHA256>"
        logger.error(error_msg)
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)
        return

    try:
        main_exe_new = Path(sys.argv[1]).resolve()
        updater_exe_new = Path(sys.argv[2]).resolve()
        version_txt_new = Path(sys.argv[3]).resolve()
        expected_sha256 = sys.argv[4]

        if not all([main_exe_new.exists(), updater_exe_new.exists(), version_txt_new.exists()]):
            raise Exception("必需文件缺失")

        if not is_admin():
            raise Exception("需要管理员权限")

        app_root = Path(sys.executable).parent
        main_exe_target = app_root / "DataCollector.exe"
        updater_exe_target = app_root / "updater.exe"
        version_txt_target = app_root / "version.txt"

        # 关闭主程序
        kill_process_by_name("DataCollector.exe")
        time.sleep(1)

        # 替换主程序
        if not verify_sha256(main_exe_new, expected_sha256):
            raise Exception("主程序校验失败")
        if main_exe_target.exists() and not force_delete(main_exe_target):
            if not handle_pending_replace(main_exe_new, main_exe_target, "主程序"):
                raise Exception("主程序替换失败")
        else:
            if not copy_with_retry(main_exe_new, main_exe_target):
                raise Exception("主程序复制失败")

        # 替换 updater
        if updater_exe_target.exists() and not force_delete(updater_exe_target):
            if not handle_pending_replace(updater_exe_new, updater_exe_target, "更新器"):
                logger.warning("更新器替换失败，将跳过")
        else:
            if not copy_with_retry(updater_exe_new, updater_exe_target):
                logger.warning("更新器复制失败，将跳过")

        # 替换 version.txt
        if not copy_with_retry(version_txt_new, version_txt_target):
            raise Exception("版本文件复制失败")

        # 清理
        temp_dir = main_exe_new.parent.parent
        if temp_dir.exists():
            shutil.rmtree(temp_dir, ignore_errors=True)

        logger.info("更新成功，重启主程序")
        restart_application(main_exe_target)

    except Exception as e:
        error_msg = f"更新失败: {str(e)}"
        logger.error(error_msg, exc_info=True)
        ctypes.windll.user32.MessageBoxW(0, error_msg, "更新失败", 0x10)


if __name__ == "__main__":
    main()